import numpy as np
from scipy.spatial import KDTree
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 1. 准备仓库位置数据 (二维坐标)
warehouses = np.array([
    [0.2, 0.3], [0.7, 0.2], [0.4, 0.7], [0.9, 0.8],
    [0.1, 0.5], [0.6, 0.9], [0.3, 0.1], [0.8, 0.4],
    [0.5, 0.6], [0.2, 0.8]
])

# 2. 构建 KD 树
kdtree = KDTree(warehouses)  # 自动选择最佳分割维度和分割点
print("KD树构建完成! 分割策略:", kdtree.tree)  # 查看树结构信息

# 3. 查询点坐标
query_point = np.array([0.4, 0.6])

# 4. 执行最近邻搜索
distance, index = kdtree.query(query_point, k=1)  # k=1 表示找最近邻
nearest_warehouse = warehouses[index]

print(f"\n查询点坐标: {query_point}")
print(f"最近仓库索引: {index} | 坐标: {nearest_warehouse}")
print(f"欧氏距离: {distance:.4f}")

# 5. 可视化结果
plt.figure(figsize=(10, 8))
plt.title("KD树最近邻搜索 - 仓库配送优化", fontsize=14)

# 绘制所有仓库
plt.scatter(
    warehouses[:, 0], warehouses[:, 1], 
    s=200, c='blue', marker='s', label='仓库位置', alpha=0.7
)

# 标记查询点
plt.scatter(
    query_point[0], query_point[1], 
    s=300, c='red', marker='*', label='查询点(待配送位置)'
)

# 标记最近邻仓库
plt.scatter(
    nearest_warehouse[0], nearest_warehouse[1], 
    s=250, c='green', marker='D', label='最近仓库'
)

# 绘制连接线
plt.plot(
    [query_point[0], nearest_warehouse[0]], 
    [query_point[1], nearest_warehouse[1]], 
    'k--', linewidth=2
)

# 添加KD树分割线 (水平/垂直线)
plt.axvline(x=0.35, color='gray', linestyle=':', alpha=0.5)  # 垂直分割线
plt.axhline(y=0.65, color='gray', linestyle=':', alpha=0.5)  # 水平分割线

plt.xlabel("经度", fontsize=12)
plt.ylabel("纬度", fontsize=12)
plt.grid(alpha=0.2)
plt.legend()
plt.show()